#!/usr/bin/env python
# -*- coding: utf-8 -*-

import cv2
import envs
import time
import rospy
import numpy as np

from img_utils import ros_to_cv2, cv2_to_ros

from std_msgs.msg import Float64
from sensor_msgs.msg import Image


class ImgObjectTheta():
    def __init__(self):
        # 图像发布者
        self.img_object_theta_pub = rospy.Publisher('/camera/image_raw/img_object_theta', Image, queue_size=1)
        # 物体斜率发布者
        self.object_theta_pub = rospy.Publisher('/camera/image_raw/object_theta', Float64, queue_size=1)
        # 图像订阅者
        self.image_sub = rospy.Subscriber('/camera/image_raw', Image, callback=self.image_callback)
        # 物体斜率
        self.object_theta = Float64()
        # canny检测相关参数
        self.ratio = 3
        self.canny_threshold = 30
        # 膨胀腐蚀相关参数
        self.kernel_size = 11
        self.iterations = 1
        # hough变换相关参数
        self.hough_threshold = 50
        self.min_hough_threshold = 0
        self.max_hough_threshold = 200
        # 颜色(BGR)
        self.color = (211, 211, 118)

        time.sleep(0.2)
        rospy.spin()

    def find_best_hough_threshold(self, img_canny):
        _min_hough_threshold = self.min_hough_threshold
        _max_hough_threshold = self.max_hough_threshold
        while True:
            if _max_hough_threshold - _min_hough_threshold <= 5:
                return _hough_threshold
            _hough_threshold = (_min_hough_threshold + _max_hough_threshold) / 2
            lines = cv2.HoughLines(img_canny, 1, np.pi/180, _hough_threshold)
            if np.all(lines == None):
                _max_hough_threshold = _hough_threshold
            else:
                if lines.shape[0] > 30:
                    _min_hough_threshold = _hough_threshold
                elif lines.shape[0] < 8:
                    _max_hough_threshold = _hough_threshold
                else:
                    return _hough_threshold

    def handle_object_theta(self, img_raw, img_canny):
        """
        物体角度
        :param img_raw: 图像原图或者压缩后的图像
        :param img_canny: 执行canny算子之后的图像
        """
        rospy.logdebug('handle object theta...')

        # hough_threshold = self.hough_threshold
        hough_threshold = self.find_best_hough_threshold(img_canny)
        # rospy.loginfo('hough_threshold: %s', hough_threshold)

        # hough变换
        lines = cv2.HoughLines(img_canny, 1, np.pi/180, hough_threshold)
        # rospy.logdebug('lines %s', lines)

        lines = lines[:, 0, :]  # 将数据转换到二维
        # 直线斜率的方差
        lines_var = np.var(lines, axis=0)
        rospy.logdebug('lines theta var is %s', lines_var[1])
        # 直线斜率的均值
        lines_mean = np.mean(lines, axis=0)
        rospy.logdebug('theta is %s', lines_mean[1])
        self.object_theta.data = lines_mean[1]
        self.object_theta_pub.publish(self.object_theta)
        # 画出所有直线
        for rho, theta in lines:
            a = np.cos(theta)
            b = np.sin(theta)
            # 从图b中可以看出x0 = rho x cos(theta) y0 = rho x sin(theta)
            x0 = a*rho
            y0 = b*rho
            # 由参数空间向实际坐标点转换
            x1 = int(x0 + 1000*(-b))
            y1 = int(y0 + 1000*a)
            x2 = int(x0 - 1000*(-b))
            y2 = int(y0 - 1000*a)
            cv2.line(img_raw, (x1, y1), (x2, y2), self.color, 2)

        self.img_object_theta_pub.publish(cv2_to_ros(img_raw, encoding='bgr8'))

        rospy.logdebug('handle object theta OK')

    def image_callback(self, data):
        rospy.logdebug('callback...')

        img = ros_to_cv2(data, encoding='bgr8')
        # canny边缘检测
        canny = cv2.Canny(img, self.canny_threshold, self.canny_threshold * self.ratio)
        # 膨胀腐蚀
        kernel = cv2.getStructuringElement(2, (self.kernel_size, self.kernel_size))
        closing = cv2.morphologyEx(canny, cv2.MORPH_CLOSE, kernel, iterations=self.iterations)
        # 第二次canny边缘检测
        canny2 = cv2.Canny(closing, 10, 30)
        self.handle_object_theta(img, canny2)


if __name__ == '__main__':
    rospy.init_node('npu_robot_img_object_theta', anonymous=True, log_level=envs.log_level)
    img_object_theta = ImgObjectTheta()
